#
import sys
sys.path.insert(0, './qlib')
from typing import Dict, Tuple, Any
import pandas as pd
import qlib
from qlib.constant import REG_CN
from qlib.utils import exists_qlib_data, init_instance_by_config, flatten_dict
from qlib.workflow import R, Recorder
from qlib.workflow.record_temp import SignalRecord, PortAnaRecord
from qlib.data import D
from apps.qte.data_source import DataSource


class QteApp(object):
    def __init__(self):
        self.name = 'apps.qte.qte_app.QteApp'

    @staticmethod
    def startup(params:Dict = {}) -> None:
        if params['task'] == 1:
            print(f'Qlib库学习 v0.0.1')
            # QteApp.init_qlib(params=params)
            QteApp.demo(params=params)
        elif params['task'] == 2:
            print(f'结构化Qlib示例程序 v0.0.2')
            QteApp.tradev1(params=params)


    @staticmethod
    def tradev1(params:Dict = {}) -> None:
        # 初始化 QLib
        print(f'Qlib初始化')
        QteApp.init_qlib()
        # 定义市场和基准
        market, benchmark = QteApp.market_bm()
        # 定义数据处理配置
        data_handler_config = QteApp.get_data_handler_config(market=market)
        task = QteApp.get_task(data_handler_config=data_handler_config)
        # 初始化模型和数据集
        model = init_instance_by_config(task["model"])
        dataset = init_instance_by_config(task["dataset"])
        rid = QteApp.train_model(model=model, dataset=dataset, task=task)
        port_analysis_config = QteApp.get_port_analysis_config(model=model, dataset=dataset, benchmark=benchmark)
        report = QteApp.run_back_test(dataset=dataset, rid=rid, port_analysis_config=port_analysis_config)
        print(f'回测完成')
        '''
        report
        report_normal_1day.pkl: <class 'pandas.core.frame.DataFrame'>; ??????
positions_normal_1day.pkl: <class 'dict'>; ??????
indicators_normal_1day.pkl: <class 'pandas.core.frame.DataFrame'>; ??????
indicators_normal_1day_obj.pkl: <class 'qlib.backtest.report.Indicator'>; ??????
port_analysis_1day.pkl: <class 'pandas.core.frame.DataFrame'>; ??????
indicator_analysis_1day.pkl: <class 'pandas.core.frame.DataFrame'>; ??????
        '''

        # r1 = report['report_normal_1day.pkl']
        # print(f'r1: {r1.tail(5)};')
        # for k, v in report['positions_normal_1day.pkl'].items():
        #     print(f'### {k}: {type(v)}; ===> {type(k)};')
        #     r3 = v
        # r2 = report['port_analysis_1day.pkl']
        # print(f'r2: {r2.tail(5)};')
        # # r3 = report['positions_normal_1day.pkl']['2020-07-31 00:00:00']
        # stock_list = r3.get_stock_list()
        # for code in stock_list:
        #     # count = r3.get_stock_count(code)
        #     price = r3.get_stock_price(code)
        #     amount = r3.get_stock_amount(code)
        #     print(f'    {code} {price} {amount};')

    @staticmethod
    def init_qlib() -> None:
        provider_uri = "./work/datas/qlib_data/cn_data"
        qlib.init(provider_uri=provider_uri, region=REG_CN)

    @staticmethod
    def market_bm() -> Tuple[str, str]:
        market = "csi300"
        benchmark = "SH000300"
        return market, benchmark

    @staticmethod
    def get_data_handler_config(market:str = 'csi300') -> Dict:
        # 定义数据处理配置
        data_handler_config = {
            "start_time": "2008-01-01",
            "end_time": "2020-08-01",
            "fit_start_time": "2008-01-01",
            "fit_end_time": "2014-12-31",
            "instruments": market
        }
        return data_handler_config

    @staticmethod
    def get_task(data_handler_config:Dict) -> Dict:
        # 定义任务配置
        task = {
            "model": {
                "class": "LGBModel",
                "module_path": "qlib.contrib.model.gbdt",
                "kwargs": {
                    "loss": "mse",
                    "colsample_bytree": 0.8879,
                    "learning_rate": 0.0421,
                    "subsample": 0.8789,
                    "lambda_l1": 205.6999,
                    "lambda_l2": 580.9768,
                    "max_depth": 8,
                    "num_leaves": 210,
                    "num_threads": 20,
                },
            },
            "dataset": {
                "class": "DatasetH",
                "module_path": "qlib.data.dataset",
                "kwargs": {
                    "handler": {
                        "class": "Alpha158",
                        "module_path": "qlib.contrib.data.handler",
                        "kwargs": data_handler_config,
                    },
                    "segments": {
                        "train": ("2008-01-01", "2014-12-31"),
                        "valid": ("2015-01-01", "2016-12-31"),
                        "test": ("2017-01-01", "2020-08-01"),
                    },
                },
            },
        }
        return task
    
    @staticmethod
    def train_model(model:Any, dataset:Any, task:Dict) -> Recorder:
        # 开始实验并训练模型
        with R.start(experiment_name="train_model"):
            R.log_params(**flatten_dict(task))
            model.fit(dataset)
            R.save_objects(trained_model=model)
            rid = R.get_recorder().id
        return rid

    @staticmethod
    def get_port_analysis_config(model:Any, dataset:Any, benchmark:str) -> Dict:
        # 定义回测配置
        port_analysis_config = {
            "executor": {
                "class": "SimulatorExecutor",
                "module_path": "qlib.backtest.executor",
                "kwargs": {
                    "time_per_step": "day",
                    "generate_portfolio_metrics": True,
                },
            },
            "strategy": {
                "class": "TopkDropoutStrategy",
                "module_path": "qlib.contrib.strategy.signal_strategy",
                "kwargs": {
                    "model": model,
                    "dataset": dataset,
                    "topk": 50,
                    "n_drop": 5,
                },
            },
            "backtest": {
                "start_time": "2017-01-01",
                "end_time": "2020-08-01",
                "account": 100000000,
                "benchmark": benchmark,
                "exchange_kwargs": {
                    "freq": "day",
                    "limit_threshold": 0.095,
                    "deal_price": "close",
                    "open_cost": 0.0005,
                    "close_cost": 0.0015,
                    "min_cost": 5,
                },
            },
        }
        return port_analysis_config

    @staticmethod
    def run_back_test(dataset:Any, rid:Recorder, port_analysis_config:Dict) -> None:
        # 执行回测和分析
        with R.start(experiment_name="backtest_analysis"):
            recorder = R.get_recorder(recorder_id=rid, experiment_name="train_model")
            model = recorder.load_object("trained_model")

            # 生成预测信号
            recorder = R.get_recorder()
            ba_rid = recorder.id
            sr = SignalRecord(model, dataset, recorder)
            sr.generate()

            # 执行回测并分析结果
            par = PortAnaRecord(recorder, port_analysis_config, "day")
            report = par.generate()
        return report



    @staticmethod
    def demo(params:Dict = {}) -> None:
        # 初始化 QLib
        provider_uri = "./work/datas/qlib_data/cn_data"
        print(f'Qlib初始化')
        qlib.init(provider_uri=provider_uri, region=REG_CN)
        # 定义市场和基准
        market = "csi300"
        benchmark = "SH000300"
        # 定义数据处理配置
        data_handler_config = {
            "start_time": "2008-01-01",
            "end_time": "2020-08-01",
            "fit_start_time": "2008-01-01",
            "fit_end_time": "2014-12-31",
            "instruments": market,
        }
        print(f'定义任务')
        # 定义任务配置
        task = {
            "model": {
                "class": "LGBModel",
                "module_path": "qlib.contrib.model.gbdt",
                "kwargs": {
                    "loss": "mse",
                    "colsample_bytree": 0.8879,
                    "learning_rate": 0.0421,
                    "subsample": 0.8789,
                    "lambda_l1": 205.6999,
                    "lambda_l2": 580.9768,
                    "max_depth": 8,
                    "num_leaves": 210,
                    "num_threads": 20,
                },
            },
            "dataset": {
                "class": "DatasetH",
                "module_path": "qlib.data.dataset",
                "kwargs": {
                    "handler": {
                        "class": "Alpha158",
                        "module_path": "qlib.contrib.data.handler",
                        "kwargs": data_handler_config,
                    },
                    "segments": {
                        "train": ("2008-01-01", "2014-12-31"),
                        "valid": ("2015-01-01", "2016-12-31"),
                        "test": ("2017-01-01", "2020-08-01"),
                    },
                },
            },
        }
        # 初始化模型和数据集
        model = init_instance_by_config(task["model"])
        dataset = init_instance_by_config(task["dataset"])
        # 开始实验并训练模型
        with R.start(experiment_name="train_model"):
            R.log_params(**flatten_dict(task))
            model.fit(dataset)
            R.save_objects(trained_model=model)
            rid = R.get_recorder().id
        # 定义回测配置
        port_analysis_config = {
            "executor": {
                "class": "SimulatorExecutor",
                "module_path": "qlib.backtest.executor",
                "kwargs": {
                    "time_per_step": "day",
                    "generate_portfolio_metrics": True,
                },
            },
            "strategy": {
                "class": "TopkDropoutStrategy",
                "module_path": "qlib.contrib.strategy.signal_strategy",
                "kwargs": {
                    "model": model,
                    "dataset": dataset,
                    "topk": 50,
                    "n_drop": 5,
                },
            },
            "backtest": {
                "start_time": "2017-01-01",
                "end_time": "2020-08-01",
                "account": 100000000,
                "benchmark": benchmark,
                "exchange_kwargs": {
                    "freq": "day",
                    "limit_threshold": 0.095,
                    "deal_price": "close",
                    "open_cost": 0.0005,
                    "close_cost": 0.0015,
                    "min_cost": 5,
                },
            },
        }
        # 执行回测和分析
        with R.start(experiment_name="backtest_analysis"):
            recorder = R.get_recorder(recorder_id=rid, experiment_name="train_model")
            model = recorder.load_object("trained_model")

            # 生成预测信号
            recorder = R.get_recorder()
            ba_rid = recorder.id
            sr = SignalRecord(model, dataset, recorder)
            sr.generate()

            # 执行回测并分析结果
            par = PortAnaRecord(recorder, port_analysis_config, "day")
            par.generate()

        print("策略回测完成！")
        


    @staticmethod
    def init_qlib(params:Dict = {}) -> None:
        # 数据存储路径
        provider_uri = "./work/datas/qlib_data/cn_data"
        # 检查数据是否存在，如果不存在则下载
        if not exists_qlib_data(provider_uri):
            print(f"Qlib 数据不存在，正在下载到 {provider_uri}")
            ds = DataSource()
            ds.qlib_data(name='qlib_data', target_dir=provider_uri, region='cn', interval='1d')
            # from qlib.tests.data import GetData
            # GetData().qlib_data(target_dir=provider_uri, region=REG_CN)
        # 初始化 QLib
        qlib.init(provider_uri=provider_uri, region=REG_CN)
        print("QLib 初始化成功！")
